/************************************************************/
/*    Cooridinate descent algorithm for models with L1      */ 
/*        parameter penalty and sparse design matrix        */
/*              Author:  Wayne Zhang                        */
/*            actuary_zhang@hotmail.com                     */
/************************************************************/

/**
 * @file slasso.c
 * @brief Cooridinate descent algorithm for lasso-like models using 
 * sparse design matrix 
 * @author Wayne Zhang                         
 */

#include "slasso.h"
#include "hlogit.h"

/** Stopping threshold in optimizing iterations */
#define LASSO_EPS 1e-7
/** Max number of iterations in the coordinate descent step */ 
#define LASSO_CD_MAXIT 10000
/** Max number of iterations in the Newton-Raphson step */ 
#define LASSO_NR_MAXIT 1000
/** Factor applied to max lambda that determines the minimum value of lambda  */
#define LASSO_LAMNRTO 0.0001
/** Factor used in modifying the probabilities  */
#define LASSO_PROB_EPS 1e-5


/************************************************/
/*           Level 1 functions                  */  
/************************************************/

/** 
 *  Soft-thresholding rules
 *
 * @param x value to be thresheld 
 * @param l the threshold
 * 
 */
static double soft(double x, double l){
  if (x > l) return(x - l);
  if (x < -l) return(x + l);
  return(0);
}


/** 
 *  Compute the scaling factor for the design matrix. 
 *  For non-intercept columns, the scaling factor is 
 *  the sample standard deviation (divided by n not n - 1)
 *  
 * @param X input sparse design matrix (compressed column)
 * @param xsd vector of the scaling factor 
 *
 */
void get_xscale(SEXP X, double *xsd){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X), xm;
  for (int j = 0; j < nB; j++){
    xsd[j] = 0;
    xm = 0;
    for (int p = xp[j]; p < xp[j + 1]; p++){      
      xm += xx[p];                             // xm := sum_i x_{ij}
      xsd[j] += xx[p] * xx[p];                 // xsd := sum_i x_{ij}^2  
    }
    xm /= nO;
    xsd[j] = sqrt(xsd[j]/nO - xm * xm);        // xsd := standard deviation (divided by n not n - 1)
    if (xsd[j] == 0) xsd[j] = 1;               // this will not be scaled (likely intercept)
  }
} 


/**
 * Compute the grid of lambdas to use
 * 
 * @param X input sparse design matrix 
 * @param y resposne
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 *
 */
void get_lambda(SEXP X, double *y, int nlam, double *lam){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X), 
    *xi = SI_SLOT(X) ;
  int nO = dim[0], nB = dim[1];   
  double *xx = SX_SLOT(X);
  double *axy = Calloc(nB, double);   
  double ym = mean(y, nO);
  
  for (int j = 0; j < nB; j++){                   // sum_i (y_i * x_{ij})
    for (int p = xp[j]; p < xp[j + 1]; p++)
      axy[j] += xx[p] * (y[xi[p]] - ym);          // why minus ybar -->> glmnet
    axy[j] = fabs(axy[j])/nO;
  }
  double lamx = dmax(axy, nB);
  double lamn = lamx * LASSO_LAMNRTO;
  double llamx = log(lamx), llamn = log(lamn); 
  double sp = (llamx - llamn)/(nlam - 1);
  for (int i = 0; i < nlam; i++)
    lam[i] = exp(llamx - i * sp);
  Free(axy);
}


/**
 * The main work horse for implementing the coordinate descent 
 * algorithm for the lasso path with sparse design matrix. The algrithm 
 * allows different weights to be used for each observation. If nlam > 1, 
 * the entire path is calculated, and the estimates are returned in the 
 * beta vector. 
 * 
 * @param X sparse design matrix 
 * @param y resposne
 * @param wts input weights (need initialization)
 * @param r vector of residuals (no initialization)
 * @beta beta vector of length nB * nlam. The first nB values are used 
 * as initial values.
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 * @param pfct the vector of penalty factors (length nB). 
 * For the intercept term, the penalty should be zero. 
 * For other terms, it is usually 1. 
 *
 */

int do_elnet(SEXP X, double *y, double *r, double *beta,  
	       double lam, double alpha, double *pfct){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X), 
    *xi = SI_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X), *b = beta;
  double *beta_old = Alloca(nB, double),
    *plam = Alloca(nB, double);
  R_CheckStack();
  double z = 0, d = 0;
  int pp = 0;

  // initialize the residual r
  Memcpy(r, y, nO);
  for (int j = 0; j < nB; j++){
    for (int p = xp[j]; p < xp[j + 1]; p++)
      r[xi[p]] -= xx[p] * beta[j];
  }  
  
  for (int j = 0; j < nB; j++) 
    plam[j] = lam * pfct[j];                                  // update threshold
    
  // run coordinate descent     
  for (int iter = 0; iter < RMM_CDITER; iter++){
    R_CheckUserInterrupt();
    Memcpy(beta_old, beta, nB);
    for (int j = 0; j < nB; j++){      
      z = 0;
      for (int p = xp[j]; p < xp[j + 1]; p++)                   // z: sum_i(x_{ij} * R_i)
	z +=  xx[p] * r[xi[p]];
      z = z/nO + beta_old[j];                                   // z:= z/nO + beta_j  
      beta[j] = TH_soft(z, plam[j] * alpha);
      beta[j] /= 1 + plam[j] * (1 - alpha);                            // Update beta_j
      d = beta_old[j] - beta[j];                                    // Update r if d != 0
      if (d != 0){                      
	for (int p = xp[j]; p < xp[j + 1]; p++)
	  r[xi[p]] += d * xx[p];
      }
    }
    
    if (dist(beta, beta_old, nB, 1) < LASSO_EPS) 
      break; 
    
  }

  if (iter == RMM_CDITER) return(1);
  return(0);
}

/**
 * Wrapper for the coordinate descent algorithm for fitting lasso paths
 * in the Gaussian model 
 * 
 * @param X input sparse design matrix, not scaled
 * @param y resposne
 * @beta beta vector of length nB * nlam. The first nB values are used 
 * as initial values.
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 * @param std binary indicator of whether the design matrix should be scaled 
 * @param pfct the vector of penalty factors which will be multiplied with lambda
 *
 */
void norm_slasso(SEXP X, double *y, double *beta, double *lam, 
		 int nlam, int std, double *pfct){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X), 
    *xi = SI_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X);
  double *wts = Calloc(nO, double), 
    *r = Calloc(nO, double), 
    *xsd = NULL; 
  for (int i = 0; i < nO; i++) wts[i] = 1.0;

  // scale the design matrix 
  if (std){
    xsd = Calloc(nB, double);
    get_xscale(X, xsd);
    for (int j = 0; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	xx[p] /= xsd[j];
    }
  }
  // get the sequence of lambdas 
  get_lambda(X, y, nlam, lam);
  
  // run lasso 
  //  do_slasso(X, y, wts, r, beta, lam, nlam, pfct);
  /*
  for (int l = 0; l < nlam; l++){
    do_slasso(X, y, wts, r, beta + l * nB, lam + l, 1, pfct);    
    // dertermine whether to truncate the lambda
    int nact = 0;
    for (int i = 0; i < nB; i++)
      if (beta[i + l * nB] != 0)
	nact++;
    if (l > 0 && nact == nB){
      if (dist(beta + l * nB, beta + (l - 1) * nB, nB, 1) < 0.0001){ 
	nlam = l;
	break; 
      }
    }
  }
  */
  // scale X and beta
  if (std){
    for (int l = 0; l < nlam; l++)
      for (int j = 0; j < nB; j++)
	beta[j + l * nB] /= xsd[j];

    for (int j = 1; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	xx[p] *= xsd[j];
    }
  }
  
  if (xsd) Free(xsd);
  Free(wts);
  Free(r);
}


/**
 * Coordinate descent algorithm for fitting lasso paths
 * in the logistic regression model 
 * 
 * @param X input sparse design matrix, not scaled
 * @param y resposne
 * @beta beta vector of length nB * nlam. The first nB values are used 
 * as initial values.
 * @param eta the vector of linear predictors
 * @param prob the vector of probabilities
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 * @param std binary indicator of whether the design matrix should be scaled 
 * @param pfct the vector of penalty factors which will be multiplied with lambda
 *
 */
/*
  static void logit_slasso(SEXP X, double *y, double *beta, double *eta, double *prob, 
  double *lam, int nlam, int std, double *pfct){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
  *xp = SP_SLOT(X), *xi = SI_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X);
  double *xsd = NULL, 
  *z = Calloc(nO, double),
  *w = Calloc(nO, double),
  *r = Calloc(nO, double),
  *beta_old = Calloc(nB, double); 

  // scale the design matrix 
  if (std){
  xsd = Calloc(nB, double);
  get_xscale(X, xsd);
  for (int j = 1; j < nB; j++){
  for (int p = xp[j]; p < xp[j + 1]; p++)
  xx[p] /= xsd[j];
  }
  }
  // get the sequence of lambdas 
  get_lambda(X, y, nlam, lam);  

  // 1): outer loop over lambdas
  for (int l = 0; l < nlam; l++){
  // 2): middle loop of the reweighted least squares 
  for (int k = 0; k < LASSO_NR_MAXIT; k++){
  Memcpy(beta_old, beta + l * nB, nB);
  // get working response and weights
  update_eta(X, beta + l * nB, eta); 
  update_prob_logit(eta, prob, nO); 
  for (int i = 0; i < nO; i++){
  // control for saturated probabilities
  if (1 - prob[i] < LASSO_PROB_EPS){
  prob[i] = 1;
  w[i] = LASSO_PROB_EPS;
  } else if (prob[i] < LASSO_PROB_EPS){
  prob[i] = 0;
  w[i] = LASSO_PROB_EPS;
  } else {
  w[i] = prob[i] * (1 - prob[i]); 
  }
  z[i] = eta[i] + (y[i] - prob[i])/w[i];
  }
  // 3): inner loop using coordinate descent
  do_slasso(X, z, w, r, beta + l * nB, lam + l, 1, pfct);
  if (dist(beta + l * nB, beta_old, nB, 1) < LASSO_EPS) 
  break; 
  if (k == LASSO_NR_MAXIT - 1)
  warning("Algorithm reaches maximum limit.");      
  }
  if (l < nlam - 1)
  Memcpy(beta + (l + 1) * nB, beta_old, nB);    
  }
  
  // scale X and beta
  if (std){
  for (int l = 0; l < nlam; l++)
  for (int j = 0; j < nB; j++)
  beta[j + l * nB] /= xsd[j];

  for (int j = 1; j < nB; j++){
  for (int p = xp[j]; p < xp[j + 1]; p++)
  xx[p] *= xsd[j];
  }
  }
  
  if (xsd) Free(xsd);
  Free(beta_old);
  Free(z);
  Free(w);
  Free(r);
  }

*/

/**
 * Coordinate descent algorithm for fitting lasso paths
 * in the multinomial logit regression model 
 * 
 * @param X input sparse design matrix, not scaled
 * @param y resposne
 * @param r the vector of residuals (length nO; no initialization)
 * @param w the vector of working weights (length nO; no initialization)
 * @param z the vector of working response (length nO; no initialization)
 * @param eta the vector of linear predictors (length nO; no initialization)
 * @param prob the vector of probabilities (length nO; no initialization)
 * @param gamma an nC * nI sparse matrix that indicates the consideration 
 * set (need initialization)
 * @param beta vector of length nB * nlam. The first nB values are used 
 * as initial values.
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 * @param std binary indicator of whether the design matrix should be scaled 
 * @param pfct the vector of penalty factors which will be multiplied with lambda
 * @param grp a vector if indices giving the starting and ending position of each 
 * group of parameters. Used in centering the betas. 
 * @param nG length of grp - 1 
 *
 */
void do_mlogit_slasso(SEXP X, double *y, double *r, double *w, double *z, 
		      double *eta, double *prob, SEXP gamma, double *beta,
		      double *lam, int nlam, double *pfct, int *grp, int nG){
  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *dim2 = SLOT_INT_NULL(gamma, "Dim"),
    *xp = SP_SLOT(X), 
    *xi = SI_SLOT(X),
    *gp = SP_SLOT(gamma),
    *gi = SI_SLOT(gamma);
  int nO = dim[0], 
    nB = dim[1], 
    nC = dim2[0], 
    nI = dim2[1],
    pp = 0; 
  double *xx = SX_SLOT(X);
  double *beta_old = Calloc(nB, double); 
  double t = 0.25, ctr = 0;
  AZERO(z, nO);

  // 1): outer loop over lambdas
  for (int l = 0; l < nlam; l++){
    // 2): middle loop of the reweighted least squares 
    for (int k = 0; k < LASSO_NR_MAXIT; k++){
      Memcpy(beta_old, beta + l * nB, nB);
      update_eta(X, beta + l * nB, eta); 
      update_prob_mlogit(gamma, eta, prob); 
      for (int i = 0; i < nO; i++)	
	  w[i] = prob[i] * (1 - prob[i]); 
      //t = 2 * dmax(w, nO);                                        // find universal lower bound of the hessian
      for (int j = 0; j < nI; j++){
	t = 2 * dmax(w + j * nC, nC);                             // individual specific bound (better)
	for (int p = gp[j]; p < gp[j + 1]; p++){
	  pp = gi[p] + j * nC;                                    // current row number for obs with nonzero prob
	  w[pp] = t;                                              // working weight
	  z[pp] = eta[pp] + (y[pp] - prob[pp])/w[pp];             // working response
	}
      }
      // 3): inner loop using coordinate descent
      do_slasso(X, z, w, r, beta + l * nB, lam + l, 1, pfct);
      
      for (int i = 0; i < nG; i++){                             // center the betas for identifiability 
	pp = grp[i + 1] - grp[i];                               // number of elments in the current group
	if (pp == 1) continue;                                  // no adjustment for alternative specific covariate
	if (i == 0)
	  ctr = mean(beta + l * nB, pp);                        // use mean for intercepts and median for others
	if (i > 0)
	  ctr = median(beta + l * nB + grp[i], pp);
	for (int j = grp[i]; j < grp[i + 1]; j++)
	  beta[l * nB + j] -= ctr;
      }      
      if (dist(beta + l * nB, beta_old, nB, 1) < LASSO_EPS) 
	break; 
      if (k == LASSO_NR_MAXIT - 1)
	warning("Newton step reaches maximum limit.");      
    }
    if (l < nlam - 1)
      Memcpy(beta + (l + 1) * nB, beta_old, nB);    
  }
}  


/**
 * Coordinate descent algorithm for fitting lasso paths
 * in the multinomial logit regression model 
 * 
 * @param X input sparse design matrix, not scaled
 * @param y resposne
 * @beta beta vector of length nB * nlam. The first nB values are used 
 * as initial values.
 * @param eta the vector of linear predictors
 * @param prob the vector of probabilities
 * @param nlam the number of lambdas
 * @param lam the vector of lambdas
 * @param std binary indicator of whether the design matrix should be scaled 
 * @param pfct the vector of penalty factors which will be multiplied with lambda
 *
 */

void mlogit_slasso(SEXP X, double *y, SEXP gamma, double *beta,
		   double *lam, int nlam, int std, double *pfct, int *grp, int nG){

  int *dim = SLOT_INT_NULL(X, "Dim"), 
    *xp = SP_SLOT(X), *xi = SI_SLOT(X);
  int nO = dim[0], nB = dim[1]; 
  double *xx = SX_SLOT(X);
  double *xsd = NULL, 
    *z = Calloc(nO, double),
    *w = Calloc(nO, double),
    *r = Calloc(nO, double),
    *eta = Calloc(nO, double),
    *prob = Calloc(nO, double),
    *beta_old = Calloc(nB, double); 
  // scale the design matrix 
  if (std){
    xsd = Calloc(nB, double);
    get_xscale(X, xsd);
    for (int j = 0; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	xx[p] /= xsd[j];
    }
  }
  // get the sequence of lambdas 
  get_lambda(X, y, nlam, lam);  

  // run mlogit lasso 
  do_mlogit_slasso(X, y, r, w, z, eta, prob, gamma, beta, 
		   lam, nlam, pfct, grp, nG);

  // scale X and beta
  if (std){
    for (int l = 0; l < nlam; l++)
      for (int j = 0; j < nB; j++)
	beta[j + l * nB] /= xsd[j];

    for (int j = 1; j < nB; j++){
      for (int p = xp[j]; p < xp[j + 1]; p++)
	xx[p] *= xsd[j];
    }
  }
  
  if (xsd) Free(xsd);
  Free(beta_old);
  Free(z);
  Free(w);
  Free(r);
  Free(eta);
  Free(prob);
}

/************************************************/
/*      Export R Callable functions             */
/************************************************/

SEXP L_update_eta(SEXP X, SEXP beta, SEXP eta){
  update_eta(X, REAL(beta), REAL(eta));
  return R_NilValue;
}

SEXP L_update_prob(SEXP gamma, SEXP eta, SEXP prob){
  update_prob_mlogit(gamma, REAL(eta), REAL(prob));
  return R_NilValue;
}



SEXP L_get_xscale(SEXP X, SEXP xsd){
  get_xscale(X, REAL(xsd));
  return R_NilValue;
}

SEXP L_get_lambda(SEXP X, SEXP y,  SEXP nlam, SEXP lam){
  get_lambda(X, REAL(y), INTEGER(nlam)[0], REAL(lam));
  return R_NilValue;
}

/*
SEXP L_update_prob_logit(SEXP eta, SEXP prob){
  update_prob_logit(REAL(eta), REAL(prob), LENGTH(eta));
  return R_NilValue;
}
*/

SEXP L_do_slasso(SEXP X, SEXP y, SEXP wts, SEXP r, SEXP beta, SEXP lam, SEXP nlam, SEXP pfct){
  do_slasso(X, REAL(y),  REAL(wts), REAL(r), REAL(beta), 
	    REAL(lam), INTEGER(nlam)[0], REAL(pfct));
  return R_NilValue;
}


SEXP L_norm_slasso(SEXP X, SEXP y, SEXP lam, SEXP nlam, SEXP std, SEXP pfct){
  int *dim = SLOT_INT_NULL(X, "Dim"); 
  int nO = dim[0], nB = dim[1];   
  SEXP beta;
  PROTECT(beta = allocMatrix(REALSXP, nB, INTEGER(nlam)[0]));
  for (int i = 0; i < nB * INTEGER(nlam)[0]; i++) REAL(beta)[i] = 0;
  norm_slasso(X, REAL(y),  REAL(beta), REAL(lam), INTEGER(nlam)[0], INTEGER(std)[0], REAL(pfct));
  UNPROTECT(1);
  return beta;
}

/*
  SEXP R_logit_slasso(SEXP X, SEXP y, SEXP eta, SEXP prob, SEXP lam, SEXP nlam, SEXP std, SEXP pfct){
  int *dim = SLOT_INT_NULL(X, "Dim"); 
  int nO = dim[0], nB = dim[1];   
  SEXP beta;
  PROTECT(beta = allocMatrix(REALSXP, nB, INTEGER(nlam)[0]));
  AZERO(REAL(beta), nB * INTEGER(nlam)[0]);
  logit_slasso(X, REAL(y),  REAL(beta),  REAL(eta), REAL(prob),
  REAL(lam), INTEGER(nlam)[0], INTEGER(std)[0], REAL(pfct));
  UNPROTECT(1);
  return beta;
  }
*/
/*
SEXP L_mlogit_slasso(SEXP X, SEXP y, SEXP gamma, SEXP lam, SEXP nlam, SEXP std, 
		     SEXP pfct, SEXP grp){
  int *dim = SLOT_INT_NULL(X, "Dim"); 
  int nO = dim[0], nB = dim[1];   
  SEXP beta;
  PROTECT(beta = allocMatrix(REALSXP, nB, INTEGER(nlam)[0]));
  AZERO(REAL(beta), nB * INTEGER(nlam)[0]);
  mlogit_slasso(X, REAL(y), gamma, REAL(beta), 
		REAL(lam), INTEGER(nlam)[0], INTEGER(std)[0], REAL(pfct),
		INTEGER(grp), LENGTH(grp) - 1);
  UNPROTECT(1);
  return beta;
}

SEXP L_do_mlogit_slasso(SEXP X, SEXP y, SEXP gamma, SEXP lam, SEXP nlam, 
			SEXP pfct, SEXP grp){
  int *dim = SLOT_INT_NULL(X, "Dim"); 
  int nO = dim[0], nB = dim[1];   
  double *xsd = NULL, 
    *z = Calloc(nO, double),
    *w = Calloc(nO, double),
    *r = Calloc(nO, double),
    *eta = Calloc(nO, double),
    *prob = Calloc(nO, double); 
  SEXP beta;
  PROTECT(beta = allocMatrix(REALSXP, nB, INTEGER(nlam)[0]));
  AZERO(REAL(beta), nB * INTEGER(nlam)[0]);
  do_mlogit_slasso(X, REAL(y), r, w, z, eta, prob, gamma, REAL(beta),
		   REAL(lam), INTEGER(nlam)[0],  
		   REAL(pfct), INTEGER(grp), LENGTH(grp) - 1);
  Free(z);
  Free(w);
  Free(r);
  Free(eta);
  Free(prob);
  UNPROTECT(1);
  return beta;
} 
*/
